
from 封装ST_GCN.train.train_model import STGCNTrainer
from 封装ST_GCN.test.test_model import STGCNTester


if __name__ == "__main__":
    train_data_path = 'C:\\Users\\Administrator\\Desktop\\Python代码\\pytorchProject1\\MY_GCN\\Data\\HanYue25_train_data.npy'
    train_label_path = 'C:\\Users\\Administrator\\Desktop\\Python代码\\pytorchProject1\\MY_GCN\\Data\\HanYue25_train_label.npy'
    test_data_path = 'C:\\Users\\Administrator\\Desktop\\Python代码\\pytorchProject1\\MY_GCN\\Data\\HanYue25_test_data.npy'
    test_label_path = 'C:\\Users\\Administrator\\Desktop\\Python代码\\pytorchProject1\\MY_GCN\\Data\\HanYue25_test_label.npy'
    num_nodes = 25
    use_residual = True
    batch_size = 32
    num_epochs = 300
    lr = 0.001
    patience = 30
    graph_args = {'layout': 'openpose', 'strategy': 'adaptive'}
    edge_importance_weighting = True

    trainer = STGCNTrainer(train_data_path, train_label_path, num_nodes, use_residual, batch_size, num_epochs, lr,
                           patience,
                           graph_args, edge_importance_weighting)
    trainer.train()

    tester = STGCNTester(test_data_path, test_label_path, num_nodes,batch_size, graph_args, edge_importance_weighting,
                         use_residual=use_residual)
    tester.load_model()
    tester.test()
